package com.llu.ssrf;

import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.InetAddress;
import java.net.URL;
import java.net.UnknownHostException;

/**
 *  防止SSRF的检查操作，但是无法防止DNS重绑定绕过
 */
public class SafeUtils {
    /**
     * check SSRF (判断逻辑为判断URL的IP是否是内网IP)
     * 如果是内网IP，返回false，表示checkSSRF不通过。否则返回true。即合法返回true
     * URL只支持HTTP协议
     * 设置了访问超时时间为3s
     */
    public static Boolean checkSSRF(String url) {
        // 连接对象
        HttpURLConnection connection;
        // 最终请求的URL地址，初始值为传入的URL
        String finalUrl = url;
        try {
            // 循环处理重定向的URL，直到成功连接或出现错误
            do {
                // 判断当前请求的URL是否是内网IP
                Boolean bRet = isInnerIpFromUrl(finalUrl);
                // 如果是内网IP，则返回false，表示处理失败
                // url - ip ，判断是否再内网
                if (bRet) {
                    return false;
                }
                // 通过当前URL打开连接
                connection = (HttpURLConnection) new URL(finalUrl).openConnection();
                // 设置不自动跳转，以便手动处理重定向
                connection.setInstanceFollowRedirects(false);
                // 设置不使用缓存，确保获取最新的响应
                connection.setUseCaches(false);
                // 设置连接超时时间为3秒
                connection.setConnectTimeout(3 * 1000);
                // 设置请求方式为GET
                connection.setRequestMethod("GET");
                // 发送连接请求
                connection.connect();
                // 获取响应码
                int responseCode = connection.getResponseCode();

                // 检查是否需要重定向
                if (responseCode >= 300 && responseCode <= 307 && responseCode != 304 && responseCode != 306) {
                    // 获取重定向的URL
                    String redirectedUrl = connection.getHeaderField("Location");
                    // 如果没有重定向URL，则停止处理
                    if (redirectedUrl == null) {
                        break;
                    }
                    // 更新最终URL为重定向的URL
                    finalUrl = redirectedUrl;
                    // 输出重定向的URL
                    System.out.println("redirected url: " + finalUrl);
                } else {
                    // 非重定向响应，停止处理
                    break;
                }
            } while (connection.getResponseCode() != HttpURLConnection.HTTP_OK);
            // 关闭连接
            connection.disconnect();
        } catch (Exception e) {
            // 在异常情况下，返回true，表示处理成功
            return true;
        }
        // 返回true，表示处理成功
        return true;

    }

    /**
     * 判断一个URL的IP是否是内网IP
     * 如果是内网IP，返回true
     * 非内网IP，返回false
     */
    public static boolean isInnerIpFromUrl(String url) throws Exception {
        String domain = getUrlDomain(url);
        if (domain.equals("")) {
            return true; // 异常URL当成内网IP等非法URL处理
        }
        //  domain --> ip DNS解析
         // test.com
        //  第一次 返回 120.26.20.xx
        //  第二次 返回 127.0.0.1

        String ip = domainToIp(domain);
        if (ip.equals("")) {
            return true; // 如果域名转换为IP异常，则认为是非法URL
        }
        return isInnerIp(ip);
    }

    /**
     * 判断给定的IP是否是内网IP
     * 内网IP范围:
     * - 10.0.0.1 - 10.255.255.254 (10.0.0.0/8)
     * - 192.168.0.1 - 192.168.255.254 (192.168.0.0/16)
     * - 127.0.0.1 - 127.255.255.254 (127.0.0.0/8)
     * - 172.16.0.1 - 172.31.255.254 (172.16.0.0/12)
     */
    public static boolean isInnerIp(String strIP) throws IOException {
        try {
            String[] ipArr = strIP.split("\\.");
            if (ipArr.length != 4) {
                return false;
            }

            int ip_split1 = Integer.parseInt(ipArr[1]);

            return (ipArr[0].equals("10") ||
                    ipArr[0].equals("127") ||
                    (ipArr[0].equals("172") && ip_split1 >= 16 && ip_split1 <= 31) ||
                    (ipArr[0].equals("192") && ipArr[1].equals("168")));
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 从URL中获取域名
     * 限制为http/https协议
     */
    public static String getUrlDomain(String url)  {
        try {
            URL u = new URL(url);
            if (!u.getProtocol().startsWith("http") && !u.getProtocol().startsWith("https")) {
                throw new IOException("Protocol error: " + u.getProtocol());
            }
            return u.getHost();
        } catch (Exception e) {
            return "";
        }
    }

    /**
     * 将域名转换为IP地址。
     * @param domain 需要转换的域名
     * @return 转换后的IP地址字符串
     */
    public static String domainToIp(String domain) {
        // 这里可以使用现有的库或者自定义逻辑来解析域名并返回其IP地址
        // 示例：使用java.net.InetAddress获取IP地址
        try {
            InetAddress ipAddress = InetAddress.getByName(domain);
            return ipAddress.getHostAddress();
        } catch (UnknownHostException e) {
            // 处理异常
            e.printStackTrace();
            return null;
        }
    }

}
